# -*- coding: utf-8 -*-
# time: 2025/3/29 10:55
# file: ch01.py
# author: hanson

import chromadb.config
# 方式一：直接创建 client 数据都存储在内存，程序运行完，数据会丢失 测试使用
client = chromadb.Client()
# 列出所有的 collection 列表：
client.list_collections()
def emb_fn():
    return [1,2,3];
collection=client.create_collection(name="my_collection")
#collection=client.create_collection(name="my_collection",embedding_function=emb_fn)


# 添加文档
collection.add(
    documents=["Article by john", "Article by Jack", "Article by Jill"],
    embeddings=[[1,2,3],[4,5,6],[7,8,9]],
    # 元数据（metadata）：是关于数据的数据或描述数据及其环境的数据，指用于描述要素、数据集或数据集系列的内容
    metadatas=[{"author": "john"}, {"author": "jack"}, {"author": "jill"}],
    ids=["1", "2", "3"])

#返回 collection 下的前 10 条数据：
data = collection.peek()
print(data)

#返回 collection 包含的数据条目数：
num = collection.count()
print(num)
# 更新文档
collection.update(
    documents=["Article by john", "Article by Jack", "Article by Jill"],
    embeddings=[[10,2,3],[40,5,6],[70,8,9]],
    metadatas=[{"author": "john"}, {"author": "jack"}, {"author": "jill"}],
    ids=["1", "2", "3"])

#查询接口
find01 = collection.get(
    ids=["1"],
    where={"author": {"$eq": "jack"}}, # 表示 metadata 中 "author" 字段值等于 "jack" 的文档
    where_document={"$contains": "john"}, # 表示文本内容中包含 "john" 的文档
)
print(find01)

#查询接口:相似文本检索接口 最重要的：：：：：：==》 检索相似性文档
query_data = collection.query(
    query_embeddings=[[1,2,3]],
    # query_texts=["Article by john"],
    n_results=3,
    where={"author": {"$eq": "john"}}, # 表示 metadata 中 "author" 字段值等于 "jack" 的文档
    where_document={"$contains": "john"}, # 表示文本内容中包含 "john" 的文档.
)
print(query_data)


# 删除 collection
client.delete_collection(name="my_collection")